package com.gouqi.gouqirpc.registry;

import com.gouqi.gouqirpc.config.RegistryConfig;
import com.gouqi.gouqirpc.model.ServiceMetaInfo;
import io.vertx.core.impl.ConcurrentHashSet;
import lombok.extern.slf4j.Slf4j;
import org.apache.curator.framework.CuratorFramework;
import org.apache.curator.framework.CuratorFrameworkFactory;
import org.apache.curator.framework.recipes.cache.CuratorCache;
import org.apache.curator.framework.recipes.cache.CuratorCacheListener;
import org.apache.curator.retry.ExponentialBackoffRetry;
import org.apache.curator.x.discovery.ServiceDiscovery;
import org.apache.curator.x.discovery.ServiceDiscoveryBuilder;
import org.apache.curator.x.discovery.ServiceInstance;
import org.apache.curator.x.discovery.details.JsonInstanceSerializer;

import java.util.Collection;
import java.util.HashSet;
import java.util.List;
import java.util.Set;
import java.util.stream.Collectors;

@Slf4j
public class ZooKeeperRegistry implements Registry {

    private CuratorFramework client;

    private ServiceDiscovery<ServiceMetaInfo> serviceDiscovery;

    /**
     * 本机注册的节点key集合（用于维护续期）
     */
    private final Set<String> localRegisterNodeKeySet = new HashSet<>();

    /**
     * 注册中心服务缓存
     */
    private final RegistryServiceCache registryServiceCache = new RegistryServiceCache();

    /**
     * 正在监听的key的集合
     */
    private final Set<String> watchKeySet = new ConcurrentHashSet<>();

    /**
     * 根节点
     */
    public static final String ZO_ROOT_PATH = "/rpc/zk";

    @Override
    public void init(RegistryConfig registryConfig) {
        client = CuratorFrameworkFactory.builder().connectString(registryConfig.getAddress()).retryPolicy(new ExponentialBackoffRetry(Math.toIntExact(registryConfig.getTimeout()), 3)).build();
        //构建serviceDiscovery
        serviceDiscovery = ServiceDiscoveryBuilder.builder(ServiceMetaInfo.class).client(client).basePath(ZO_ROOT_PATH).serializer(new JsonInstanceSerializer<>(ServiceMetaInfo.class)).build();
        try {
            //启动client和serviceDiscovery
            client.start();
            serviceDiscovery.start();
        } catch (Exception e) {
            e.printStackTrace();
            throw new RuntimeException(e);
        }
    }

    @Override
    public void register(ServiceMetaInfo serviceMetaInfo) throws Exception {
        //注册到Zookeeper
        serviceDiscovery.registerService(buildServiceInstance(serviceMetaInfo));

        //设置要存储的键值对
        String registryKey = ZO_ROOT_PATH + "/" + serviceMetaInfo.getServiceNodeKey();

        //添加节点信息到本地缓存
        localRegisterNodeKeySet.add(registryKey);
    }

    @Override
    public void unRegister(ServiceMetaInfo serviceMetaInfo) {
        try {
            serviceDiscovery.unregisterService(buildServiceInstance(serviceMetaInfo));
        } catch (Exception e) {
            e.printStackTrace();
            throw new RuntimeException(e);
        }
        String registryKey = ZO_ROOT_PATH + "/" + serviceMetaInfo.getServiceNodeKey();
        //移除缓存
        localRegisterNodeKeySet.remove(registryKey);
    }

    @Override
    public List<ServiceMetaInfo> serviceDiscovery(String serviceKey) {
        //优先从缓存获取服务
        List<ServiceMetaInfo> cacheServiceMetaInfoList = registryServiceCache.readCache();
        if (cacheServiceMetaInfoList != null) {
            return cacheServiceMetaInfoList;
        }
        try {
            //查询服务信息
            Collection<ServiceInstance<ServiceMetaInfo>> serviceInstances = serviceDiscovery.queryForInstances(serviceKey);

            //解析服务信息
            List<ServiceMetaInfo> serviceMetaInfoList = serviceInstances.stream().map(ServiceInstance::getPayload).collect(Collectors.toList());

            //写入服务缓存
            registryServiceCache.writeCache(serviceMetaInfoList);
            return serviceMetaInfoList;

        } catch (Exception e) {
            e.printStackTrace();
            throw new RuntimeException("获取服务列表失败", e);
        }

    }

    @Override
    public void destroy() {
        //下线节点
        System.out.println("当前节点下线");
        //遍历本节点下的所有key，下线节点，因为都是临时节点，服务下线就会被删掉
        for (String key : localRegisterNodeKeySet) {
            try {
                client.delete().guaranteed().forPath(key);
            } catch (Exception e) {
                e.printStackTrace();
                throw new RuntimeException(key + "节点下线失败");
            }
        }
        //释放资源
        if (client != null) {
            client.close();
        }
    }

    @Override
    public void heartBeat() {
        //不需要心跳检测机制，建立临时节点，如果服务器故障，临时节点直接丢失
    }

    @Override
    public void watch(String serviceNodeKey) {

        String watchKey = ZO_ROOT_PATH + "/" + serviceNodeKey;
        //之前未被监听，开启监听
        boolean newWatch = watchKeySet.add(watchKey);
        if (newWatch) {
            CuratorCache curatorCache = CuratorCache.build(client, watchKey);
            curatorCache.start();
            curatorCache.listenable().addListener(
                    CuratorCacheListener.builder()
                            .forDeletes(childData -> {
                                registryServiceCache.clearCache();
                            }).forChanges(((oldNode, node) ->
                                    registryServiceCache.clearCache()
                            )).build()
            );
        }
    }

    private ServiceInstance<ServiceMetaInfo> buildServiceInstance(ServiceMetaInfo serviceMetaInfo) {
        String serviceAddress = serviceMetaInfo.getServiceHost() + ":" + serviceMetaInfo.getServicePort();
        try {
            return ServiceInstance
                    .<ServiceMetaInfo>builder()
                    .id(serviceAddress)
                    .name(serviceMetaInfo.getServiceKey())
                    .address(serviceAddress)
                    .payload(serviceMetaInfo)
                    .build();
        } catch (Exception e) {
            e.printStackTrace();
            throw new RuntimeException(e);
        }
    }

}
